// The MPC formulation for r^2, m^2 and h, similar to the one in Ni Dong's dissertation.
// One decision variable, T(t_k)

#include "MPC_EW2D_PatternedW_c.h"
#include <sstream>
static int FUNC_EVAL_COUNTER = 0;
// -----------------------------------------------------------------------
/* Constructor. */
MPC1_PatternedW::MPC1_PatternedW(int i_mode, int i_LatticeSize,
				 double i_XYmax, double i_Freq,double i_t_end,
				 double i_r2_set, double i_q_r,
				 double i_m2_set, double i_q_m):
MPC_Formulation(4,10),
mode(i_mode), LatticeSize(i_LatticeSize),
XYmax(i_XYmax), Freq(i_Freq),t_end(i_t_end),
r2_set(i_r2_set),q_r(i_q_r/(i_r2_set*i_r2_set)),
m2_set(i_m2_set),q_m(i_q_m/(i_m2_set*i_m2_set))
{
//  int m,n;

  // Input parameters ---------------
  W_ini = 0.03;    // unit: nm/s
  A_ini = 0.01;
  W_old = W_ini;
  A_old = A_ini;
//  Freq  = 5;
  T     = 460;

  // Output parameters -?------------
  num_decision_val = 2;

  // Cost function -----------------
//  r2_set = 30*30;
//  m2_set = 0.2*0.2;
//  q_r    = 1.0/(r2_set*r2_set);
//  q_m    = 1.0/(m2_set*m2_set);

  // Constraints -------------------
  W_min = 0.02;
  W_max = 0.4;
//  dW_max = 0.01;
  dW_max = 1.0;
  A_min = 0.0;
  A_max = 2e19; // positive infinity

  // Model parameters --------------
//  LatticeSize = 100;
//  XYmax = 8000;
//  t_end = 100.0;
  dt    = 1.0;
//  mode  = 50;

  int p;
  for(p=0;p<4;p++){
    cov_z_m[p]  = new double[(1+mode)*(1+mode)];
    mean_z_m[p] = new double[(1+mode)*(1+mode)];
  }

  double dx = XYmax/LatticeSize;

  // Parameterize the optimization problem
  cost_fun = new eq_f(r2_set,q_r,m2_set,q_m,t_end,mode,LatticeSize,XYmax,Freq);
  
  /*
  cost_fun->parameterize(0,r2_set);
  cost_fun->parameterize(1,q_r);
  cost_fun->parameterize(2,m2_set);
  cost_fun->parameterize(3,q_m);
  cost_fun->parameterize(4,mode);
  cost_fun->parameterize(5,LatticeSize);
  cost_fun->parameterize(6,XYmax);
  cost_fun->parameterize(7,t_end);
  cost_fun->parameterize(8,Freq);
  */

  cost_fun->parameterize(0,cov_z_m[0]);
  cost_fun->parameterize(1,mean_z_m[0]);
  cost_fun->parameterize(2,cov_z_m[1]);
  cost_fun->parameterize(3,mean_z_m[1]);
  cost_fun->parameterize(4,cov_z_m[2]);
  cost_fun->parameterize(5,mean_z_m[2]);
  cost_fun->parameterize(6,cov_z_m[3]);
  cost_fun->parameterize(7,mean_z_m[3]);

  reset();
}
// -----------------------------------------------------------------------
MPC1_PatternedW::~MPC1_PatternedW(){
  for(int p=0;p<4;p++){
    delete []cov_z_m[p];
    delete []mean_z_m[p];
  }
  delete cost_fun;
}
// -----------------------------------------------------------------------
void MPC1_PatternedW::parameterize(double CurrentTime, double** inputs){
  int    p,m,n;
  int    index;
  double dt;
  for(m=0;m<=mode;m++){
    for(n=0;n<=mode;n++){
      index = index2D(m,n);
      for(p=0;p<4;p++){
	mean_z_m[p][index] = inputs[p][index];
	if(p==2 && m == Freq && n == 0)
	  cov_z_m[p][index] = 0.0;
	else
	  // cov_z_m[p][index]  = mean_z_m[p][index]*mean_z_m[p][index]; // FIXME
	  cov_z_m[p][index] = 0.0;
      }	
    }
  }
  time = CurrentTime;
  dt = t_end - CurrentTime+MPC1_PatternedW::dt;

  cost_fun->parameterize(7,dt);
}
// -----------------------------------------------------------------------
void MPC1_PatternedW::reset(){
  W_old = W_ini;
  A_old = A_ini;
}
// -----------------------------------------------------------------------
void MPC1_PatternedW::setPara(int para_id,double para_val){
  // Put those tunable parameters here. For those un-tunable ones,
  // hard-code them in a derived class.
  assert(para_id < MPC_Formulation::num_para);
  switch(para_id){
  case 0:
    r2_set = para_val;
    cost_fun->parameterize(0,r2_set);
    break;
  case 1:
    q_r = para_val/(r2_set*r2_set);
    cost_fun->parameterize(1,q_r);
    break;
  case 2:
    m2_set = para_val;
    cost_fun->parameterize(2,m2_set);
    break;
  case 3:
    q_m = para_val/(m2_set*m2_set);
    cost_fun->parameterize(3,q_m);
    break;
  case 4:
    Freq = para_val;
    cost_fun->parameterize(8,Freq);
    break;
  case 5:
    T = para_val;
    break;
  case 6:
    LatticeSize = int(para_val);
    cost_fun->parameterize(5,LatticeSize);
    break;
  case 7:
    XYmax = para_val;
    cost_fun->parameterize(6,XYmax);
    break;
  case 8:
    mode = int(para_val);
    cost_fun->parameterize(4,mode);
    break;
  case 9:
    dt = para_val;
    cost_fun->parameterize(7,dt);
    break;
  default:
    cerr << "MPC1_PatternedW::setPara(): Invalid para_id (=" << para_id << "\n";
    exit(-1);
  }
}
// -----------------------------------------------------------------------
void MPC1_PatternedW::get_solution(double *x){
  x[0] = T;
  x[1] = W_old;
  x[2] = A_old;
  x[3] = Freq;
}
// -----------------------------------------------------------------------


// -----------------------------------------------------------------------
//                 Optimization Formulation
// -----------------------------------------------------------------------
bool MPC1_PatternedW::get_nlp_info(Index& n, Index& m, Index& nnz_jac_g,
				   Index& nnz_h_lag, IndexStyleEnum& index_style)
{
  // The problem described in MPC1_PatternedW.hpp has 3 variables,
  //  x0 = W0, x1 = A
  n = 2;

  // 1 inequality constraints: 
  // 0 <= g(x) = W-A = x[1]-x[2] <= \infty
  m = 1;
  // The following limits on variables do not consider as inequality constraints
  // in IPOPT terminology
  // max{T_{min},T_old-dT_max} <= T <= min{T_{max},T_old+dT_max}
  // max{W_{min},W_old-dW_max} <= W <= min{W_{max},W_old+dW_max}
  // 0 <= A <= infty

  // 2 nonzeros in the jacobian (all elements are non-zero)
  nnz_jac_g = 2;

  // and 2 nonzeros in the hessian of the lagrangian
  // (one in the hessian of the objective for x2,
  //  and one in the hessian of the constraints for x1)
  nnz_h_lag = 3;

  // We use the standard fortran index style for row/col entries
  index_style = C_STYLE;

  return true;
}
// -----------------------------------------------------------------------
bool MPC1_PatternedW::get_bounds_info(Index n, Number* x_l, Number* x_u,
				      Index m, Number* g_l, Number* g_u)
{
  // here, the n and m we gave IPOPT in get_nlp_info are passed back to us.
  // If desired, we could assert to make sure they are what we think they are.
  assert(n == 2);
  assert(m == 1);

  // x0 = W0
  x_l[0] = max(W_min,W_old-dW_max);
  x_u[0] = min(W_max,W_old+dW_max);

  // x2 = A
  x_l[1] = A_min;
  x_u[1] = A_max;

  // we have one equality constraint, so we set the bounds on this constraint
  // to be equal (and zero).
  g_l[0] = 0.0;
  g_u[0] = 2e19;

  return true;
}
// -----------------------------------------------------------------------
bool MPC1_PatternedW::get_starting_point(Index n, bool init_x, Number* x,
					 bool init_z, Number* z_L, Number* z_U,
					 Index m, bool init_lambda, Number* lambda)
{
  // Here, we assume we only have starting values for x, if you code
  // your own NLP, you can provide starting values for the others if
  // you wish.
  assert(init_x == true);
  assert(init_z == false);
  assert(init_lambda == false);

  // we initialize x in bounds, in the upper right quadrant
  x[0] = W_old;
  x[1] = A_old;

// #ifdef DEBUG
  double x_min[2];
  double x_max[2];
  double dx[2];
  double W_lb = max(W_min,W_old-dW_max);
  double W_ub = min(W_max,W_old+dW_max);
  dx[0] = 0.01*(W_ub-W_lb);
  dx[1] = 0.01*(W_ub);
  x_min[0] = W_lb;
  x_max[0] = W_ub;
  x_min[1] = 0;
  x_max[1] = W_ub;
  
  /*
  stringstream ss;
  ss << FUNC_EVAL_COUNTER++;
  string filename("cost_fun");
  filename = filename+ss.str()+".dat";
  cost_fun->plot(x_min,x_max,dx,filename);
  */
// #endif
  return true;
}
// -----------------------------------------------------------------------
bool MPC1_PatternedW::eval_f(Index n, const Number* x, bool new_x, 
			     Number& obj_value)
{
  // return the value of the objective function
  // J = q_m*(m2_sp-m2_f)^2+q_r*(r2_sp-r2_f)^2+q_h*(h_sp-h_f)^2

  obj_value = cost_fun->value(x);
  //  double epsilon[2]={1e-8,1e-8};
  //  try{
  //    MathFunction* f = dynamic_cast<eq_f*>(&cost_fun);
  //    perturb_MathFunction(f,x,epsilon);
  //  } catch(exception& e){cout << e.what();}
  
  return true;
}
// -----------------------------------------------------------------------
bool MPC1_PatternedW::eval_grad_f(Index n, const Number* x, bool new_x, 
				  Number* grad_f)
{
  // return the gradient of the objective function grad_{x} f(x)
  cost_fun->gradient(x,grad_f);
  return true;
}
// -----------------------------------------------------------------------
bool MPC1_PatternedW::eval_g(Index n, const Number* x, bool new_x, Index m, 
			     Number* g){
  // return the value of the constraints: g(x)
  Number W = x[0];
  Number A = x[1];

  g[0] = W-A;

  return true;
}
// -----------------------------------------------------------------------
bool MPC1_PatternedW::eval_jac_g(Index n, const Number* x, bool new_x,
				 Index m, Index nele_jac, 
				 Index* iRow, Index *jCol,Number* values)
{
  if (values == NULL) {
    // return the structure of the jacobian of the constraints

    // element at 1,1: grad_{x1} g_{1}(x)
    iRow[0] = 0;    jCol[0] = 0;

    // element at 1,2: grad_{x2} g_{1}(x)
    iRow[1] = 0;    jCol[1] = 1;
  }
  else {
    // return the values of the jacobian of the constraints
    // element at 1,1: grad_{x1} g_{1}(x)
    values[0] = 1;

    // element at 1,2: grad_{x1} g_{1}(x)
    values[1] = -1;
  }

  return true;
}
// -----------------------------------------------------------------------
bool MPC1_PatternedW::eval_h(Index n, const Number* x, bool new_x,
			     Number obj_factor, 
			     Index m, const Number* lambda,bool new_lambda, 
			     Index nele_hess,
			     Index* iRow, Index* jCol, Number* values)
{
  // The hessian matrix is a 3-by-3 matrix
  if (values == NULL) {
    // return the structure. This is a symmetric matrix, fill the lower left
    // triangle only.
    
    iRow[0] = 0;    jCol[0] = 0;
    iRow[1] = 1;    jCol[1] = 0;
    iRow[2] = 1;    jCol[2] = 1;

  }
  else {
    // return the values
    // fill the objective portion
    values[0] = 0;
    values[1] = 0;
    values[2] = 0;
    
    // add the portion for the first constraint, which is all zero, so do NOTHING here.
  }
  return false; // USE Quasi-Newton method
}
// -----------------------------------------------------------------------
void MPC1_PatternedW::finalize_solution(SolverReturn status,
					Index n, const Number* x, 
					const Number* z_L, const Number* z_U,
					Index m, const Number* g, const Number* lambda,
					Number obj_value,
					const IpoptData* ip_data,
					IpoptCalculatedQuantities* ip_cq)
{
  // Save the current solution, which will be used as the initial solution.

  // If time == t_end, the result is discard, keep the latest result.
  if(time!=t_end){
    W_old = x[0];
    A_old = x[1];
  }

//  double r2_f;
//  double m2_f;

  //  eq_predict_state(T_old,W_old,A_old,Freq);
  //  r2_f = eq_r2_f();
  //  m2_f = eq_m2_f();

  //  printf("*** time = %f\nT = %f, W = %f, A = %f\nr_f = %f, m2_f = %f\n",
  // time,x[0],x[1],x[2],r2_f,m2_f);
}
// -----------------------------------------------------------------------

